/*
 * EquilibriumProblem.cpp
 *
 *  Created on: 27 Jul 2011
 *      Author: Allan
 */

#include "EquilibriumProblem.h"

// GeoReact includes
#include "EquilibriumSolver.h"

EquilibriumProblem::EquilibriumProblem()
{}

EquilibriumProblem::EquilibriumProblem(const Multiphase& multiphase, const CanonicalReactionSystem& reactions) :
multiphase(multiphase), reactions(reactions)
{
	// The primary (jSpecies) and equilibrium (jSpecies) species
	const vector<string> jSpecies = reactions.GetPrimarySpecies();
	const vector<string> eSpecies = reactions.GetSecondarySpecies();
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Nj = jSpecies.size();
	Ne = eSpecies.size();
	
	// Set the indexes of the primary (jIndexes) and equilibrium (eIndexes) species
	jIndexes = multiphase[jSpecies];
	eIndexes = multiphase[eSpecies];
	
	// Set the canonical stoichiometric matrix of the system of reactions
	vej = reactions.GetCanonicalStoichiometricMatrix();
	
	// Allocate all the auxiliary vector/matrix variables with zero entries
	Ke = VectorXd::Zero(Ne);	Qe = VectorXd::Zero(Ne);
	nj = VectorXd::Zero(Nj);	ne = VectorXd::Zero(Ne);
	aj = VectorXd::Zero(Nj);	ae = VectorXd::Zero(Ne);
	h  = VectorXd::Zero(Nj);
	
	alpha = MatrixXd::Zero(Nj, Nj);
	beta  = MatrixXd::Zero(Nj, Ne);
	
	hfunctions.resize(Nj); 
}

EquilibriumProblem::~EquilibriumProblem()
{}

void EquilibriumProblem::SetEquilibriumConditions(const vector<EquilibriumCondition>& eConditions)
{
	// Loop over all the equilibrium conditions
	for(unsigned i = 0; i < eConditions.size(); ++i)
	{
		// Set the i-th row of the alpha matrix
		alpha.row(i) = eConditions[i].alphaRow;
		
		// Set the i-th row of the beta matrix
		beta.row(i) = eConditions[i].betaRow;
		
		// Set the i-th row of the equilibrium condition functions
		hfunctions[i] = eConditions[i].h;
	}
}

void EquilibriumProblem::SetParameters(double T, double P, const VectorXd& n)
{
	this->T = T;
	this->P = P;
	this->n = n;
	
	// Initialise the equilibrium constants of the equilibrium reactions (Ke)
	Ke = reactions.CanonicalEquilibriumConstants(T, P);
}

void EquilibriumProblem::Function(const VectorXd& njne, VectorXd& F)
{
	// Update the auxiliary data members of the equilibrium problem
	UpdateAuxiliaryData(njne);
	
	// Compute the residual vector Fj = F[0:Nj-1]
	F.segment(00,  Nj) = (alpha * nj) + (beta * ne) - h;
	
	// Compute the residual vector Fe = F[Nj:Ne-1]
	F.segment(Nj,  Ne) = (Ke.array() / Qe.array()).log();
}

void EquilibriumProblem::Jacobian(const VectorXd& njne, MatrixXd& J)
{
	// Compute the block jacobian matrix Jjj = J[0:Nj-1, 0:Nj-1]
	J.block(00, 00, Nj, Nj) = alpha;
	
	// Compute the block jacobian matrix Jje = J[0:Nj-1, Nj:Ne-1]
	J.block(00, Nj, Nj, Ne) = beta;
	
	// Compute the block jacobian matrix Jej = J[Nj:Ne-1, 0:Nj-1]
	J.block(Nj, 00, Ne, Nj) = - vej * nj.asDiagonal().inverse();
	
	// Compute the block jacobian matrix Jej = J[Nj:Ne-1, Nj:Ne-1]
	J.block(Nj, Nj, Ne, Ne) = ne.asDiagonal().inverse();
}

void EquilibriumProblem::UpdateAuxiliaryData(const VectorXd& njne)
{
	// Check if all the entries in (njne) are non-negative
	if(njne.minCoeff() > 0)
	{
		// Update the nj and ne part of the vector (n) with the values of the vector (njne)
		for(unsigned j = 0; j < Nj; ++j) n[jIndexes[j]] = njne[j];
		for(unsigned e = 0; e < Ne; ++e) n[eIndexes[e]] = njne[e + Nj];
		
		// Update the activity vector (a) with temperature (t), pressure (P) and the mole vector (n)
		a = multiphase.Activities(T, P, n);
		
		// Update the reaction quotients of the equilibrium reactions
		Qe = reactions.CanonicalReactionQuotients(a);
		
		// Update the vector (h), the right-hand side values of the speciation conditions
		for(unsigned j = 0; j < Nj; ++j) h[j] = hfunctions[j](T, P, n, a);
	}
	
	// Update the mole vectors (nj) and (ne)
	for(unsigned j = 0; j < Nj; ++j) nj[j] = njne[j];
	for(unsigned e = 0; e < Ne; ++e) ne[e] = njne[e + Nj];
}
